import optuna
import matplotlib.pyplot as plt
import numpy as np
import multiprocessing
from optuna.samplers import TPESampler
from concurrent.futures import ProcessPoolExecutor
from coilgun import coilgunApi
from coilgun import figure_show
from coilgun import CoilMode
import copy  
import time

optuna.logging.set_verbosity(optuna.logging.WARNING)


class ai_Optuna:
    
    class SharedMemoryRecorder:
        """共享内存记录器"""
        def __init__(self, manager):
            self.params = manager.list()
            self.scores = manager.list()
            self.constraints = manager.list()
        
        def add_trial(self, params, score, constraint):
            self.params.append(params)
            self.scores.append(score)
            self.constraints.append(constraint)
        
        def get_data(self):
            return {
                "params": list(self.params),
                "scores": list(self.scores),
                "constraints": list(self.constraints)
            }

    def __init__(self, lin_preset,CONSTRAINT_MAX,TOTAL_TRIALS,NUM_PROCESSES):
        self.start = time.perf_counter() 


        self.lin_preset =lin_preset
        self.only_id = self.lin_preset['only_id']

        self.PARAM_BOUNDS = self.lin_preset['PARAM_BOUNDS']

        #print(self.PARAM_BOUNDS)

        # ====== 配置参数 ======
        
        ## 电压弹长扫描模式 采用穷举计算 ，否则用 ai优化模式
        '''TOTAL_TRIALS_ture = 1
        for param in self.PARAM_BOUNDS  :
            print(len(param[3]))
            if param[2]== "tuple" :  
                TOTAL_TRIALS_ture *= len(param[3])
            else:
                TOTAL_TRIALS_ture =0
                break
        print(f'TOTAL_TRIALS_ture={TOTAL_TRIALS_ture}')    

        if TOTAL_TRIALS_ture ==0:
            self.CONSTRAINT_MAX = CONSTRAINT_MAX          # 必须小于的约束值
            self.TOTAL_TRIALS = TOTAL_TRIALS
            self.NUM_PROCESSES = NUM_PROCESSES
        else:
            self.CONSTRAINT_MAX = CONSTRAINT_MAX          # 必须小于的约束值
            self.TOTAL_TRIALS = TOTAL_TRIALS_ture
            self.NUM_PROCESSES = 1'''

        self.CONSTRAINT_MAX = CONSTRAINT_MAX          # 必须小于的约束值
        self.TOTAL_TRIALS = TOTAL_TRIALS
        self.NUM_PROCESSES = NUM_PROCESSES
        
        if self.PARAM_BOUNDS !=[]:
            self.main()
        else:
            print('没有可优化参数')



    def perset_assignment(self,params):
        for i,PARAM_B,param in zip (range(len(params)), self.PARAM_BOUNDS ,  params   ) :
            name = eval(PARAM_B[0])
            #print(f'i:{i},PARAM_B:{PARAM_B},param:{param}  ')
            if name[0] != -1 : #参数化更新线圈参数
                self.lin_preset["coil_layout"][name[0]][name[1]]= param
            else:             #参数化更新非线圈参数
                self.lin_preset[name[1]]  = param

        return self.lin_preset


    def objective(self,trial, recorder, constraint_max, param_bounds):
        """目标函数"""
        '''params = [
            trial.suggest_float(name, low, high) if param_type=='float' else trial.suggest_int(name, low, high)
            for name,param_type, low, high in param_bounds
        ]'''
        params =[]

        for name,param_type,param_type2, data in param_bounds :
            if param_type == 'float' and param_type2 == 'list':
                params.append(trial.suggest_float(name, data[0] , data[1]  ) )

            elif param_type == 'int' and param_type2 == 'list':
                params.append(trial.suggest_int(name, data[0] , data[1]  ))

            elif param_type2 == 'tuple':
                params.append(trial.suggest_categorical(name, data  ))



        
    
        # ====== 自定义计算部分 ======

        nowperset2= copy.deepcopy( self.perset_assignment(params) )

        #print(params)
        #计算
        score,constraint ,eff , v_result, m_All   = coilgunApi(nowperset2)

        # ====== 结束自定义 ======
        
        recorder.add_trial(params, score, constraint)
        return score if constraint < constraint_max else float('-inf')

    def run_optimization(self,recorder, constraint_max, param_bounds, n_trials):

        study = optuna.create_study(
            direction="maximize",
            sampler=TPESampler(n_startup_trials=50)
        )
        study.optimize(
            lambda trial: self.objective(trial, recorder, constraint_max, param_bounds),
            n_trials=n_trials ,
        )
        return recorder.get_data()

    def visualize_results(self,data, constraint_max):
        """交互式可视化函数"""
        fig, ax = plt.subplots(figsize=(12, 8))


        # 准备数据
        params = np.array(data["params"])
        scores = np.array(data["scores"])
        constraints = np.array(data["constraints"])
        
        # 绘制散点图
        sc = ax.scatter(scores, constraints, c='blue', alpha=0.6, picker=True)
        ax.axhline(constraint_max, color='red', linestyle='--')
        
        # 初始化标注
        annot = ax.annotate("", xy=(0,0), xytext=(20,20),
                            textcoords="offset points",
                            bbox=dict(boxstyle="round", fc="w"),
                            arrowprops=dict(arrowstyle="->"))
        annot.set_visible(False)
        
        # 存储数据参考
        fig.data = {
            "params": params,
            "scores": scores,
            "constraints": constraints
        }
        fig.current_idx = None
        fig.scatter = sc
        
        def update_annot(ind):
            """更新标注内容"""
            idx = ind["ind"][0]
            pos = sc.get_offsets()[idx]
            annot.xy = pos
            text = f"Params: {params[idx]}\nScore: {scores[idx]:.2f}\nConstraint: {constraints[idx]:.2f}"
            annot.set_text(text)
            annot.get_bbox_patch().set_alpha(0.8)
            
            # 标红选中点
            colors = ['blue'] * len(scores)
            colors[idx] = 'red'
            sc.set_color(colors)
            fig.canvas.draw_idle()

        def hover(event):
            """鼠标悬停事件处理"""
            vis = annot.get_visible()
            if event.inaxes == ax:
                cont, ind = sc.contains(event)
                if cont:
                    update_annot(ind)
                    annot.set_visible(True)
                    fig.current_idx = ind["ind"][0]
                    fig.canvas.draw_idle()
                else:
                    if vis:
                        annot.set_visible(False)
                        # 恢复颜色
                        sc.set_color('blue')
                        fig.canvas.draw_idle()

        def on_right_click(event):
            """右键点击事件处理"""
            if event.button == 3 and fig.current_idx is not None:
                idx = fig.current_idx

 
                nowperset2= copy.deepcopy( self.perset_assignment(params[idx]) )


                ###not_scanning_main(  CoilMode(0, nowperset2, f' {self.only_id } | {params[idx]}'          ) )     gg

                nowperset2["only_id"] = f' {self.only_id } | {params[idx]}'
                resultdata = coilgunApi(nowperset2,'非参数扫描模式')
                figure_show(resultdata[6] , resultdata[7] , resultdata[8] , resultdata[9], resultdata[10]   )


                # 创建新窗口显示详细信息



        # 绑定事件
        fig.canvas.mpl_connect("motion_notify_event", hover)
        fig.canvas.mpl_connect("button_press_event", on_right_click)
        
        # 图表装饰
        ax.set_xlabel("增加的动能 (J)")
        ax.set_ylabel("最大电流 (A)")
        ax.set_title(f"{self.only_id}\n(右键点击标红点查看详细信息)")
        plt.tight_layout()
        plt.show()



    def main(self):

        
        # 初始化共享内存
        with multiprocessing.Manager() as manager:
            # 创建主记录器
            main_recorder = self.SharedMemoryRecorder(manager)
            
            # 创建进程池
            with ProcessPoolExecutor(max_workers=self.NUM_PROCESSES) as executor:
                futures = []


                
                trials_per_process = self.TOTAL_TRIALS // self.NUM_PROCESSES
                
                # 每个进程使用独立记录器
                for _ in range(self.NUM_PROCESSES):
                    process_recorder = self.SharedMemoryRecorder(manager)
                    future = executor.submit(
                        self.run_optimization,
                        process_recorder,
                        self.CONSTRAINT_MAX,
                        self.PARAM_BOUNDS,
                        trials_per_process
                    )
                    futures.append(future)
                
                # 收集所有数据
                all_data = [f.result() for f in futures]
            
            # 合并数据
            merged_data = {
                "params": [p for d in all_data for p in d["params"]],
                "scores": [s for d in all_data for s in d["scores"]],
                "constraints": [c for d in all_data for c in d["constraints"]]
            }
            
            print( f"调用次数：{len(merged_data['params'])  }"  )
            end = time.perf_counter()
            print(f"耗时: {end - self.start:.6f} 秒") 

            # 显示最优解
            feasible_scores = [s if c < self.CONSTRAINT_MAX else -1  for s, c in zip(merged_data["scores"], merged_data["constraints"])  ]
            if feasible_scores:
                best_idx = np.argmax(feasible_scores)
                print(f"\n=== 最优解 ===")
                print(f"参数: {merged_data['params'][best_idx]}")
                print(f"增加的动能: {merged_data['scores'][best_idx]:.2f}")
                print(f"最大电流: {merged_data['constraints'][best_idx]:.2f}")
            else:
                print("未找到可行解")

            # 可视化
            self.visualize_results(merged_data, self.CONSTRAINT_MAX)
            




if __name__ == "__main__":
    multiprocessing.set_start_method("spawn")
    startOptuna=ai_Optuna()